class UnsqueezeDimension(object):
    """Add a dimension at position 0."""
    def __call__(self, sample):
        return sample.unsqueeze(0)

class SelectElectrodes(object):
    def __init__(self, num_electrodes):
        self.num_electrodes = num_electrodes

    def __call__(self, sample):
        return sample[:self.num_electrodes]